
from langchain.chains import RetrievalQA
from langchain_openai import AzureChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型
from tool import get_azure_endpoint,get_api_version,get_api_key
if __name__ == '__main__':
    llm = AzureChatOpenAI(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        azure_deployment="gpt-4o-mini",  # 重命名为 azure_deployment
        openai_api_version=get_api_version(),  # 参数名不变
        openai_api_key=get_api_key(),
        openai_api_type="azure",
    )
    embedding = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
    )
    vectordb_chinese = Chroma(
        persist_directory="./docs/chroma/matplotlib",
        embedding_function=embedding,
    )
    # 声明一个检索式问答链
    qa_chain = RetrievalQA.from_chain_type(
    llm,
        retriever=vectordb_chinese.as_retriever()
    )
    # 可以以该方式进行检索问答
    question = "这节课的主要话题是什么"
    result = qa_chain({"query": question})
    print(result["result"])